import asyncio
from concurrent.futures import ThreadPoolExecutor
import websockets
import json
import base64
import _thread as thread
from sse_starlette.sse import EventSourceResponse
from typing import AsyncGenerator
from dataclasses import dataclass
from websocket import WebSocketApp
import websocket
import hashlib
import hmac
from urllib.parse import urlencode
from datetime import datetime
from time import mktime
from wsgiref.handlers import format_date_time
import ssl
import time
import threading
from asyncio import Queue
from asr_service import SparkASRService
STATUS_FIRST_FRAME = 0
STATUS_CONTINUE_FRAME = 1
STATUS_LAST_FRAME = 2


@dataclass
class ASRConfig:
    app_id: str
    api_key: str
    api_secret: str
    host: str = "ws-api.xfyun.cn"
    path: str = "/v2/iat"


@dataclass
class Segment:
    score: int = None
    text: str = None


@dataclass
class Vad:
    begin: int = None
    end: int = None


@dataclass
class Transcation:
    index: int
    segment: Segment = None


@dataclass
class ASRResult:
    transcation: Transcation=None
    vad: Vad = None
    begin: int = None
    end: int = None
    sid: str = None
    bestMatchText: str = None
    status: int = None
    full_text: str = None


@dataclass
class ASRError:
    sid: str = None
    code: int = None
    errMsg: str = None


class ASRCallbacks:
    def onResult(result: ASRResult, usrTag: str):
        print("onResult: ", result)

    def onError(error: ASRError, usrTag: str):
        print("onError: ", error)


@dataclass
class AudioAttributes:
    sampleRate: int = 16000
    encoding: str = "raw"  # speex(8k),speex-wb(16k),lame(mp3,仅支持中英文)
    channels: int = 1
    bitepth: int = 16
    frame_size = 0
    # 16k音频：audio/L16;rate=16000
    # 8k音频：audio/L16;rate=8000


class ASRServer:
    config: ASRConfig = None
    # audio_buffers = {}  # client_id -> audio buffer
    # asr_connections = {}  # client_id -> ASR WebSocket
    executor = ThreadPoolExecutor(max_workers=100)

    def __init__(self, client_id,frame_size:int=8000):
        if not client_id:
            self.client_id = id(self)
        else:
            self.client_id = client_id
        self.ws = None
        self.audio_buffer = bytearray()
        self.callbacks = ASRCallbacks()
        self.full_text = ""
        self.queue = Queue()
        self.run = True
        self.buffer_lock = threading.Lock()
        self.loop = asyncio.get_event_loop()
        self.frame_size = frame_size

    @classmethod
    def create_asr_server(cls, config: ASRConfig = None, client_id: int = None,frame_size:int=8000):
        if config is None and cls.config is None:
            raise ValueError("ASRConfig is None")
        elif cls.config is None:
            cls.config = config

        server = ASRServer(client_id,frame_size)
        return server

    def _create_url(self):
        url = 'wss://ws-api.xfyun.cn/v2/iat'
        # url = 'wss://iat-api.xfyun.cn/v2/iat'
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = "host: " + "ws-api.xfyun.cn" + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + "/v2/iat " + "HTTP/1.1"
        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(ASRServer.config.api_secret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha = base64.b64encode(
            signature_sha).decode(encoding='utf-8')

        authorization_origin = "api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"" % (
            ASRServer.config.api_key, "hmac-sha256", "host date request-line", signature_sha)
        authorization = base64.b64encode(
            authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": "ws-api.xfyun.cn"
        }
        # 拼接鉴权参数，生成url
        url = url + '?' + urlencode(v)
        # print("date: ",date)
        # print("v: ",v)
        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，比对相同参数时生成的url与自己代码生成的url是否一致
        # print('websocket url :', url)
        return url

    def create_asr_connection(self):
        """Create WebSocket connection to ASR service"""
        def on_message(ws, message):
            try:
                data = json.loads(message)

                if data["code"] != 0:
                    print(f"ASR error: {data['message']}")
                    return
                result: ASRResult = ASRResult()
                result.sid = data["sid"]
                result.status = data["data"]["status"]
                # result.transcation.index = data["data"]["result"]["sn"]
                # result.begin = data["data"]["result"]["bg"]
                # result.end = data["data"]["result"]["ed"]

                bestMatchText = ""
                for item in data["data"]["result"]["ws"]:
                    for w in item["cw"]:
                        bestMatchText += w["w"]
                print("bestMatchText:",bestMatchText)
                result.bestMatchText = bestMatchText
                result.full_text = self.full_text + bestMatchText
                self.full_text = result.full_text
                asyncio.run_coroutine_threadsafe(self.queue.put(result), self.loop)
                print("get result:",self.full_text)
            except Exception as e:
                print(f"Error processing ASR message: {e}")

        def on_error(ws, error):
            print(f"ASR connection error: {error}")

        def on_close(ws, *args):
            print("ASR connection closed")
            self.run=False
            if self.ws:
                del self.ws

        def on_open(ws):
            intervel = 0.04  # 发送音频间隔(单位:s)

            def run(*args):
                # Send first frame with config
                d = {
                    "common": {"app_id": ASRServer.config.app_id},
                    "business": {
                        "domain": "iat",
                        "language": "zh_cn",
                        "accent": "mandarin",
                        "vinfo": 1,
                        "vad_eos": 3000
                    },
                    "data": {
                        "status": STATUS_FIRST_FRAME,
                        "format": "audio/L16;rate=16000",
                        "encoding": "raw"
                    }
                }
                ws.send(json.dumps(d))
                # print("send first frame")

                # Continuously send audio data
                last_send_time = time.time()
                while self.run:
                    try:
                        # if len(self.audio_buffer) > 0:

                        with self.buffer_lock:
                            print("buffer size: ", len(self.audio_buffer))
                            if len(self.audio_buffer) >= self.frame_size:
                                chunk = self.audio_buffer[:self.frame_size]
                                del self.audio_buffer[:self.frame_size]
                                print(f"send:{len(chunk)}")
                                status = STATUS_CONTINUE_FRAME
                                if len(self.audio_buffer) == 0:
                                    status = STATUS_LAST_FRAME

                                data = {
                                    "data": {
                                        "status": status,
                                        "format": "audio/L16;rate=16000",
                                        "audio": chunk.decode(),
                                        "encoding": "raw"
                                    }
                                }
                                ws.send(json.dumps(data))
                                last_send_time = time.time()
                                # print("send...")

                        # Timeout if no data for 500ms
                        if time.time() - last_send_time > 0.4:
                            with self.buffer_lock:
                                if len(self.audio_buffer) > 0:
                                    # print(
                                    #     f"timeout last frame:{len(self.audio_buffer)}")
                                    chunk = self.audio_buffer[:]
                                    print(f"send-all:{len(chunk)}")
                                    del self.audio_buffer[:]

                            data = {
                                "data": {
                                    "status": STATUS_LAST_FRAME,
                                    "format": "audio/L16;rate=16000",
                                    "audio": chunk.decode(),
                                    "encoding": "raw"
                                }
                            }
                            ws.send(json.dumps(data))
                            break
                        elif len(self.audio_buffer) <= 0:
                            if time.time() - last_send_time > 1:
                                break

                        time.sleep(intervel)
                        print("buffer size: ", len(self.audio_buffer))
                    except Exception as e:
                        print(f"Error in ASR send loop: {e}")
                        break
                print("out running....")

            ASRServer.executor.submit(run)

        url = self._create_url()
        self.ws = websocket.WebSocketApp(
            url,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close,
            on_open=on_open
        )

    async def write_audio(self, audio_data: bytes, base64str: str = None):
        if audio_data:
            base64str = base64.b64encode(audio_data).decode()
            base64bytes = base64str.encode('utf-8')
            # base64bytes = audio_data
        else:
            base64bytes = base64str.encode('utf-8')

        with self.buffer_lock:
            if not self.audio_buffer:
                self.audio_buffer = bytearray(base64bytes)
            else:
                self.audio_buffer.extend(base64bytes)
                print(f"Audio buffer updated: {len(base64bytes)} bytes")

        if not self.ws.sock or not self.ws.sock.connected:
            ASRServer.executor.submit(
                self.ws.run_forever,sslopt={
                    "cert_reqs": ssl.CERT_NONE,
                    "check_hostname": False
                })
            await asyncio.sleep(0.1)

    def stop(self):
        self.run = False
        if self.ws:
            self.ws.close()
            del self.ws
            print("ASR connection closed by stop")
            self.audio_buffer = None
            self.callbacks = None

    async def process_results(self, callback=None) -> AsyncGenerator[ASRResult,None]:
        """Async generator to process results from the queue"""
        print("Starting ASR result processing")
        
        # 初始等待队列中有数据或超时
        first_wait = True
        
        while True:
            try:
                # 如果是第一次等待且队列为空，等待一段时间
                if first_wait and self.queue.empty():
                    try:
                        result = await asyncio.wait_for(self.queue.get(), timeout=1.0)
                    except asyncio.TimeoutError:
                        if not self.run:
                            break
                        continue
                    first_wait = False
                else:
                    if self.queue.empty() and not self.run:
                        break
                    result = await self.queue.get()
                
                print("Processing ASR result")
                if callback:
                    await callback(result)
                yield result
                
                # 处理完当前结果后，如果队列为空且run=false则退出
                if self.queue.empty() and not self.run:
                    break
                    
            except Exception as e:
                print(f"Error processing result: {e}")
                break
                
        print("Finished processing all results")
